"""
不使用 QStandardItem 的折叠表格实现
使用 QAbstractTableModel，直接管理数据，性能更好
"""

from PySide6.QtWidgets import (
    QApplication, QMainWindow, QTableView, QVBoxLayout, 
    QWidget, QPushButton, QHBoxLayout, QStyledItemDelegate, QStyle
)
from PySide6.QtGui import QColor, QFont, QPainter, QPixmap, QIcon
from PySide6.QtCore import Qt, QRect, QSize, QAbstractTableModel, QModelIndex
import sys


class CollapseDelegate(QStyledItemDelegate):
    """折叠按钮代理"""
    
    def __init__(self, parent=None):
        super().__init__(parent)
        self.collapsed_icon = self._create_collapsed_icon()
        self.expanded_icon = self._create_expanded_icon()
    
    def _create_collapsed_icon(self):
        """创建折叠图标（右箭头 ▶）"""
        pixmap = QPixmap(16, 16)
        pixmap.fill(Qt.GlobalColor.transparent)
        painter = QPainter(pixmap)
        painter.setRenderHint(QPainter.RenderHint.Antialiasing)
        
        from PySide6.QtGui import QPolygon, QBrush
        from PySide6.QtCore import QPoint
        triangle = QPolygon([
            QPoint(5, 4), QPoint(5, 12), QPoint(11, 8)
        ])
        painter.setBrush(QBrush(QColor(100, 100, 100)))
        painter.setPen(Qt.PenStyle.NoPen)
        painter.drawPolygon(triangle)
        painter.end()
        return QIcon(pixmap)
    
    def _create_expanded_icon(self):
        """创建展开图标（下箭头 ▼）"""
        pixmap = QPixmap(16, 16)
        pixmap.fill(Qt.GlobalColor.transparent)
        painter = QPainter(pixmap)
        painter.setRenderHint(QPainter.RenderHint.Antialiasing)
        
        from PySide6.QtGui import QPolygon, QBrush
        from PySide6.QtCore import QPoint
        triangle = QPolygon([
            QPoint(4, 5), QPoint(12, 5), QPoint(8, 11)
        ])
        painter.setBrush(QBrush(QColor(100, 100, 100)))
        painter.setPen(Qt.PenStyle.NoPen)
        painter.drawPolygon(triangle)
        painter.end()
        return QIcon(pixmap)
    
    def paint(self, painter, option, index):
        """绘制单元格"""
        if index.column() != 1:  # 第1列（名称列）
            super().paint(painter, option, index)
            return
        
        model = index.model()
        is_group = model.is_group_row(index.row())
        
        painter.save()
        
        # 绘制背景
        if option.state & QStyle.StateFlag.State_Selected:
            painter.fillRect(option.rect, option.palette.highlight())
        else:
            bg_color = index.data(Qt.ItemDataRole.BackgroundRole)
            if bg_color:
                painter.fillRect(option.rect, bg_color)
            else:
                painter.fillRect(option.rect, option.palette.base())
        
        if is_group:
            # 绘制折叠/展开图标
            is_expanded = model.is_group_expanded(index.row())
            icon = self.expanded_icon if is_expanded else self.collapsed_icon
            
            icon_rect = QRect(option.rect.x() + 5, 
                            option.rect.y() + (option.rect.height() - 16) // 2, 16, 16)
            icon.paint(painter, icon_rect)
            
            # 绘制文本（粗体）
            text_rect = option.rect.adjusted(25, 0, 0, 0)
            painter.setPen(option.palette.text().color())
            font = QFont()
            font.setBold(True)
            painter.setFont(font)
            text = index.data(Qt.ItemDataRole.DisplayRole)
            painter.drawText(text_rect, Qt.AlignmentFlag.AlignLeft | Qt.AlignmentFlag.AlignVCenter, text)
        else:
            # 普通数据行，缩进显示
            text_rect = option.rect.adjusted(35, 0, 0, 0)
            painter.setPen(option.palette.text().color())
            text = index.data(Qt.ItemDataRole.DisplayRole)
            painter.drawText(text_rect, Qt.AlignmentFlag.AlignLeft | Qt.AlignmentFlag.AlignVCenter, text)
        
        painter.restore()
    
    def sizeHint(self, option, index):
        size = super().sizeHint(option, index)
        return QSize(size.width(), max(size.height(), 28))


class CollapseTableModel(QAbstractTableModel):
    """折叠表格模型 - 不使用 QStandardItem"""
    
    def __init__(self, parent=None):
        super().__init__(parent)
        self._headers = ['No', '名称', '数量', '单价', '金额', '备注']  # 添加 No 列
        self._all_rows = []  # 所有行数据 [{'type': 'group'/'data', 'data': [...], 'expanded': True/False, ...}]
        self._visible_rows = []  # 当前可见的行索引
        self._groups = {}  # 分组信息 {group_name: group_index}
        self._data_row_counter = 0  # 数据行计数器
    
    def rowCount(self, parent=QModelIndex()):
        return len(self._visible_rows)
    
    def columnCount(self, parent=QModelIndex()):
        return len(self._headers)
    
    def data(self, index, role=Qt.ItemDataRole.DisplayRole):
        if not index.isValid():
            return None
        
        visible_row = index.row()
        if visible_row >= len(self._visible_rows):
            return None
        
        actual_row = self._visible_rows[visible_row]
        row_data = self._all_rows[actual_row]
        
        if role == Qt.ItemDataRole.DisplayRole or role == Qt.ItemDataRole.EditRole:
            return row_data['data'][index.column()]
        
        elif role == Qt.ItemDataRole.BackgroundRole:
            if row_data['type'] == 'group':
                return QColor(230, 240, 250)
            return None
        
        elif role == Qt.ItemDataRole.TextAlignmentRole:
            if row_data['type'] == 'data' and index.column() > 0:
                return Qt.AlignmentFlag.AlignCenter
            return Qt.AlignmentFlag.AlignLeft | Qt.AlignmentFlag.AlignVCenter
        
        return None
    
    def headerData(self, section, orientation, role=Qt.ItemDataRole.DisplayRole):
        if orientation == Qt.Orientation.Horizontal and role == Qt.ItemDataRole.DisplayRole:
            return self._headers[section]
        return None
    
    def is_group_row(self, visible_row):
        """判断是否是分组行"""
        if visible_row >= len(self._visible_rows):
            return False
        actual_row = self._visible_rows[visible_row]
        return self._all_rows[actual_row]['type'] == 'group'
    
    def is_group_expanded(self, visible_row):
        """获取分组的展开状态"""
        if visible_row >= len(self._visible_rows):
            return False
        actual_row = self._visible_rows[visible_row]
        row_data = self._all_rows[actual_row]
        return row_data.get('expanded', True)
    
    def toggle_group(self, visible_row):
        """切换分组的展开/折叠状态"""
        if visible_row >= len(self._visible_rows):
            return
        
        actual_row = self._visible_rows[visible_row]
        row_data = self._all_rows[actual_row]
        
        if row_data['type'] != 'group':
            return
        
        # 切换状态
        row_data['expanded'] = not row_data['expanded']
        
        # 重建可见行列表
        self._rebuild_visible_rows()
        
        # 通知视图刷新
        self.layoutChanged.emit()
    
    def _rebuild_visible_rows(self):
        """重建可见行列表"""
        self._visible_rows.clear()
        
        for i, row_data in enumerate(self._all_rows):
            if row_data['type'] == 'group':
                # 分组行始终可见
                self._visible_rows.append(i)
            elif row_data['type'] == 'data':
                # 数据行：检查所属分组是否展开
                parent_group = row_data.get('parent_group')
                if parent_group is not None:
                    group_row = self._all_rows[parent_group]
                    if group_row.get('expanded', True):
                        self._visible_rows.append(i)
    
    def add_group(self, group_name, children_data):
        """添加分组及其子数据"""
        # 添加分组行
        group_index = len(self._all_rows)
        self._data_row_counter += 1
        self._all_rows.append({
            'type': 'group',
            'data': [str(self._data_row_counter), group_name, '', '', '', ''],  # 第0列为空，第1列为分组名
            'expanded': True,
            'name': group_name
        })
        self._groups[group_name] = group_index
        
        # 添加子数据行
        for child_data in children_data:
            self._data_row_counter += 1
            # 第0列是行号
            row_data = [str(self._data_row_counter)] + child_data
            self._all_rows.append({
                'type': 'data',
                'data': row_data,
                'parent_group': group_index
            })
        
        # 重建可见行
        self._rebuild_visible_rows()
        self.layoutChanged.emit()
    
    def expand_all(self):
        """展开所有分组"""
        for row_data in self._all_rows:
            if row_data['type'] == 'group':
                row_data['expanded'] = True
        self._rebuild_visible_rows()
        self.layoutChanged.emit()
    
    def collapse_all(self):
        """折叠所有分组"""
        for row_data in self._all_rows:
            if row_data['type'] == 'group':
                row_data['expanded'] = False
        self._rebuild_visible_rows()
        self.layoutChanged.emit()


class CollapseTableView(QTableView):
    """支持折叠功能的表格视图"""
    
    def __init__(self, parent=None):
        super().__init__(parent)
        
        # 隐藏表头
        self.horizontalHeader().hide()
        self.verticalHeader().hide()
        
        # 设置代理
        self.collapse_delegate = CollapseDelegate()
        self.setItemDelegateForColumn(1, self.collapse_delegate)  # 第1列（名称）
    
    def mousePressEvent(self, event):
        """处理点击事件"""
        if event.button() == Qt.MouseButton.LeftButton:
            index = self.indexAt(event.pos())
            if index.isValid() and index.column() == 1:  # 第1列（名称）
                model = self.model()
                if model.is_group_row(index.row()):
                    model.toggle_group(index.row())
                    event.accept()
                    return
        super().mousePressEvent(event)


class CollapseTableDemo(QMainWindow):
    """演示主窗口"""
    
    def __init__(self):
        super().__init__()
        self.init_ui()
        self.setup_model()
        self.setup_data()
    
    def init_ui(self):
        """初始化UI"""
        self.setWindowTitle('折叠表格示例（无QStandardItem）')
        self.setGeometry(100, 100, 800, 600)
        
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout(central_widget)
        
        # 按钮工具栏
        toolbar = QHBoxLayout()
        
        self.btn_expand_all = QPushButton('展开全部')
        self.btn_expand_all.clicked.connect(lambda: self.model.expand_all())
        toolbar.addWidget(self.btn_expand_all)
        
        self.btn_collapse_all = QPushButton('折叠全部')
        self.btn_collapse_all.clicked.connect(lambda: self.model.collapse_all())
        toolbar.addWidget(self.btn_collapse_all)
        
        self.btn_add_group = QPushButton('添加分组')
        self.btn_add_group.clicked.connect(self.add_group)
        toolbar.addWidget(self.btn_add_group)
        
        toolbar.addStretch()
        layout.addLayout(toolbar)
        
        # TableView
        self.table_view = CollapseTableView()
        self.table_view.setAlternatingRowColors(True)
        layout.addWidget(self.table_view)
    
    def setup_model(self):
        """设置模型"""
        self.model = CollapseTableModel()
        self.table_view.setModel(self.model)
    
    def setup_data(self):
        """设置示例数据"""
        groups = [
            {
                'name': '电子产品',
                'children': [
                    ['笔记本电脑', '5', '5000', '25000', '畅销'],
                    ['平板电脑', '10', '3000', '30000', '热门'],
                    ['手机', '20', '4000', '80000', '新品'],
                ]
            },
            {
                'name': '家具',
                'children': [
                    ['办公桌', '15', '800', '12000', '促销'],
                    ['办公椅', '25', '500', '12500', ''],
                ]
            },
            {
                'name': '文具',
                'children': [
                    ['笔记本', '100', '10', '1000', ''],
                    ['钢笔', '50', '20', '1000', ''],
                    ['文件夹', '80', '5', '400', ''],
                ]
            }
        ]
        
        for group in groups:
            self.model.add_group(group['name'], group['children'])
        
        self.table_view.resizeColumnsToContents()
        self.statusBar().showMessage('已加载3个分组（使用QAbstractTableModel，无QStandardItem）')
    
    def add_group(self):
        """添加新分组"""
        import random
        group_name = f'新分组{random.randint(1, 100)}'
        children_data = [
            [f'商品{i}', str(random.randint(1, 50)), str(random.randint(10, 1000)), '', '']
            for i in range(3)
        ]
        
        self.model.add_group(group_name, children_data)
        self.table_view.resizeColumnsToContents()
        self.statusBar().showMessage(f'已添加分组：{group_name}')


def main():
    """主函数"""
    app = QApplication(sys.argv)
    window = CollapseTableDemo()
    window.show()
    sys.exit(app.exec())


if __name__ == '__main__':
    main()
